#!/usr/bin/env python3
"""
LaTeX Title Translation Script

This script translates chapter, section, subsection, and subsubsection titles 
from English to Chinese in LaTeX documents, specifically for technical documentation
about C++ function call tree analysis and MCP integration.

Features:
- Comprehensive technical term dictionary
- Backup functionality
- Dry-run mode
- Preserves LaTeX formatting
- Handles escaped characters properly
- Smart detection of already translated content
"""

import os
import re
import shutil
import argparse
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Tuple, Set

class LaTeXTitleTranslator:
    def __init__(self):
        self.translation_dict = self._build_translation_dictionary()
        self.processed_files = []
        self.changes_made = []
        
    def _build_translation_dictionary(self) -> Dict[str, str]:
        """Build comprehensive technical translation dictionary."""
        return {
            # Document Structure & Academic Terms
            "Introduction and Software Overview": "引言与软件概述",
            "Architecture Design": "架构设计", 
            "Parsing Engines": "解析引擎",
            "Analysis Algorithms": "分析算法",
            "Validation Framework": "验证框架",
            "Model Context Protocol Integration": "模型上下文协议集成",
            "Conclusions and Future Work": "结论与未来工作",
            
            # System Architecture
            "System Architecture Overview": "系统架构概述",
            "Architectural Principles": "架构原则",
            "Component Design": "组件设计",
            "Core Components": "核心组件",
            "Design Patterns": "设计模式",
            "Data Flow Architecture": "数据流架构",
            "Analysis Pipeline": "分析流水线",
            "Data Structures": "数据结构",
            "Scalability and Performance": "可扩展性与性能",
            "Performance Optimization Strategies": "性能优化策略",
            "Scalability Architecture": "可扩展性架构",
            "Error Handling and Resilience": "错误处理与弹性",
            "Error Classification": "错误分类",
            "Resilience Mechanisms": "弹性机制",
            
            # Parsing & Analysis
            "Overview of Parsing Strategies": "解析策略概述",
            "Parsing Engine Selection Criteria": "解析引擎选择标准",
            "Regex-Based Parsing Engine": "基于正则表达式的解析引擎",
            "Pattern-Based Analysis Approach": "基于模式的分析方法",
            "Core Pattern Definitions": "核心模式定义",
            "Preprocessing and Normalization": "预处理与标准化",
            "Limitations and Constraints": "局限性与约束",
            "Clang AST Parsing Engine": "Clang AST解析引擎",
            "Abstract Syntax Tree Analysis": "抽象语法树分析",
            "LibClang Integration": "LibClang集成",
            "Parser Coordination and Selection": "解析器协调与选择",
            "Hybrid Parsing Strategy": "混合解析策略",
            "Result Integration": "结果集成",
            "Performance Characteristics": "性能特征",
            "Comparative Analysis": "比较分析",
            "Optimization Strategies": "优化策略",
            
            # Algorithms
            "Call Graph Construction": "调用图构建",
            "Graph Theoretical Foundations": "图论基础",
            "Construction Algorithm": "构建算法",
            "Name Resolution and Disambiguation": "名称解析与消歧",
            "Cycle Detection Algorithms": "循环检测算法",
            "Strongly Connected Components": "强连通组件",
            "Cycle Classification": "循环分类",
            "Tree Generation and Visualization": "树生成与可视化",
            "Tree Extraction from Graphs": "图中树提取",
            "Multi-Root Tree Generation": "多根树生成",
            "Complexity Analysis": "复杂度分析",
            "Algorithmic Complexity": "算法复杂度",
            "Statistical Analysis": "统计分析",
            "Graph Metrics": "图指标",
            "Distribution Analysis": "分布分析",
            
            # Validation
            "Validation Architecture Overview": "验证架构概述",
            "Validation Philosophy": "验证理念",
            "Validation Architecture Diagram": "验证架构图",
            "Confidence Scoring System": "置信度评分系统",
            "Confidence Model": "置信度模型",
            "Confidence Categories": "置信度类别",
            "Call Relationship Validation": "调用关系验证",
            "Consistency Checking": "一致性检查",
            "Semantic Validation": "语义验证",
            "Cycle Detection and Analysis": "循环检测与分析",
            "Cycle Validation Framework": "循环验证框架",
            "Cycle Examples": "循环示例",
            "Numerical Relationship Analysis": "数值关系分析",
            "Complexity Metrics Validation": "复杂度指标验证",
            "Runtime Cross-Validation": "运行时交叉验证",
            "Profiling Data Integration": "分析数据集成",
            "Discrepancy Analysis": "差异分析",
            "Validation Reporting": "验证报告",
            "Comprehensive Validation Reports": "综合验证报告",
            "Actionable Recommendations": "可执行建议",
            
            # MCP Integration
            "MCP Architecture Overview": "MCP架构概述",
            "Protocol Foundations": "协议基础",
            "MCP Architecture Diagram": "MCP架构图",
            "Tool Implementation": "工具实现",
            "Core Analysis Tools": "核心分析工具",
            "Parameter Schema Definition": "参数模式定义",
            "Asynchronous Processing Model": "异步处理模型",
            "Data Serialization and Transport": "数据序列化与传输",
            "JSON Serialization Strategy": "JSON序列化策略",
            "Streaming Data Transfer": "流式数据传输",
            "Client Integration Patterns": "客户端集成模式",
            "Development Environment Integration": "开发环境集成",
            "Usage Patterns and Best Practices": "使用模式与最佳实践",
            "Analysis Workflow Diagram": "分析工作流程图",
            "Performance and Scalability": "性能与可扩展性",
            "Resource Management": "资源管理",
            
            # Solution Approach
            "Solution Approach": "解决方案",
            "Design Philosophy": "设计理念",
            "Technical Approach": "技术方法",
            "Key Innovations": "关键创新",
            "Document Structure": "文档结构",
            
            # Common Technical Terms
            "Implementation": "实现",
            "Framework": "框架",
            "System": "系统",
            "Analysis": "分析",
            "Algorithm": "算法",
            "Strategy": "策略",
            "Methodology": "方法论",
            "Approach": "方法",
            "Overview": "概述",
            "Introduction": "引言",
            "Conclusion": "结论",
            "Summary": "总结",
            "Evaluation": "评估",
            "Assessment": "评价",
            "Validation": "验证",
            "Integration": "集成",
            "Configuration": "配置",
            "Deployment": "部署",
            "Architecture": "架构",
            "Design": "设计",
            "Development": "开发",
            "Testing": "测试",
            "Debugging": "调试",
            "Optimization": "优化",
            "Performance": "性能",
            "Scalability": "可扩展性",
            "Security": "安全",
            "Reliability": "可靠性",
            "Maintenance": "维护",
            "Documentation": "文档",
            
            # Specific compound terms that need special handling
            "Confidence Weighting": "置信度权重",
            "Conflict Resolution": "冲突解决",
            "Complementary Analysis": "补充分析",
            "Caching": "缓存",
            "Incremental Processing": "增量处理",
            "Parallel Processing": "并行处理",
            "Cross-Validation": "交叉验证",
            "Load Balancing": "负载均衡",
            "Distributed Analysis": "分布式分析",
            "Cache Strategy": "缓存策略",
            "Connection Pooling": "连接池化",
            "Memory Management": "内存管理",
            "CPU Throttling": "CPU节流",
            "Disk Space Management": "磁盘空间管理",
        }
    
    def _is_chinese_content(self, text: str) -> bool:
        """Check if text contains significant Chinese content."""
        chinese_chars = re.findall(r'[\u4e00-\u9fff]', text)
        return len(chinese_chars) > len(text) * 0.3
    
    def _translate_title(self, title: str) -> str:
        """Translate a single title using the dictionary."""
        # If already contains Chinese, likely already translated
        if self._is_chinese_content(title):
            return title
            
        # Direct lookup first
        if title in self.translation_dict:
            return self.translation_dict[title]
        
        # Try partial matching for compound titles
        translated_parts = []
        remaining_text = title
        
        # Sort by length (longer matches first)
        sorted_terms = sorted(self.translation_dict.keys(), key=len, reverse=True)
        
        for term in sorted_terms:
            if term in remaining_text:
                translated_parts.append((term, self.translation_dict[term]))
                remaining_text = remaining_text.replace(term, "<<<REPLACED>>>")
        
        if translated_parts:
            result = title
            for original, translated in translated_parts:
                result = result.replace(original, translated)
            return result
        
        # If no translation found, return original
        return title
    
    def _process_latex_command(self, match) -> str:
        """Process a single LaTeX sectioning command."""
        command = match.group(1)  # chapter, section, etc.
        title = match.group(2)    # the title content
        
        # Clean up any LaTeX escaping in the title
        clean_title = title.replace('\\ ', ' ').strip()
        translated_title = self._translate_title(clean_title)
        
        # If translation was made, record it
        if translated_title != clean_title:
            self.changes_made.append((clean_title, translated_title))
        
        # Reconstruct the command with translated title
        return f"\\{command}{{{translated_title}}}"
    
    def translate_file(self, file_path: str, dry_run: bool = False) -> bool:
        """
        Translate titles in a single LaTeX file.
        Returns True if changes were made.
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
        except Exception as e:
            print(f"Error reading {file_path}: {e}")
            return False
        
        # Pattern to match LaTeX sectioning commands
        # Matches: \chapter{...}, \section{...}, \subsection{...}, \subsubsection{...}
        pattern = r'\\(chapter|section|subsection|subsubsection)\{([^}]+)\}'
        
        original_content = content
        self.changes_made = []  # Reset for this file
        
        # Apply translations
        new_content = re.sub(pattern, self._process_latex_command, content)
        
        # Check if any changes were made
        if new_content != original_content:
            if not dry_run:
                # Write the translated content back
                try:
                    with open(file_path, 'w', encoding='utf-8') as file:
                        file.write(new_content)
                    print(f"✓ Translated {len(self.changes_made)} titles in {file_path}")
                    for original, translated in self.changes_made:
                        print(f"    '{original}' → '{translated}'")
                except Exception as e:
                    print(f"Error writing {file_path}: {e}")
                    return False
            else:
                print(f"[DRY RUN] Would translate {len(self.changes_made)} titles in {file_path}:")
                for original, translated in self.changes_made:
                    print(f"    '{original}' → '{translated}'")
            
            self.processed_files.append(file_path)
            return True
        else:
            print(f"- No translations needed for {file_path}")
            return False
    
    def create_backup(self, base_path: str) -> str:
        """Create a backup of the entire directory."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_path = f"{base_path}_backup_{timestamp}"
        
        try:
            shutil.copytree(base_path, backup_path)
            print(f"✓ Created backup at: {backup_path}")
            return backup_path
        except Exception as e:
            print(f"Error creating backup: {e}")
            raise
    
    def find_tex_files(self, directory: str) -> List[str]:
        """Find all .tex files in the directory and subdirectories."""
        tex_files = []
        
        for root, dirs, files in os.walk(directory):
            for file in files:
                if file.endswith('.tex'):
                    tex_files.append(os.path.join(root, file))
        
        return sorted(tex_files)
    
    def translate_directory(self, directory: str, dry_run: bool = False, 
                          create_backup: bool = True) -> Dict[str, int]:
        """
        Translate all .tex files in a directory.
        Returns a summary of changes made.
        """
        if not os.path.isdir(directory):
            raise ValueError(f"Directory does not exist: {directory}")
        
        # Create backup unless it's a dry run
        if create_backup and not dry_run:
            self.create_backup(directory)
        
        tex_files = self.find_tex_files(directory)
        if not tex_files:
            print("No .tex files found.")
            return {"files_processed": 0, "files_changed": 0, "total_translations": 0}
        
        print(f"Found {len(tex_files)} .tex files to process...")
        
        files_changed = 0
        total_translations = 0
        
        for tex_file in tex_files:
            print(f"\nProcessing: {tex_file}")
            if self.translate_file(tex_file, dry_run):
                files_changed += 1
                total_translations += len(self.changes_made)
        
        summary = {
            "files_processed": len(tex_files),
            "files_changed": files_changed,
            "total_translations": total_translations
        }
        
        return summary

def main():
    parser = argparse.ArgumentParser(
        description="Translate LaTeX chapter, section, and subsection titles from English to Chinese"
    )
    parser.add_argument(
        "directory", 
        help="Directory containing .tex files to translate"
    )
    parser.add_argument(
        "--dry-run", "-d",
        action="store_true",
        help="Preview changes without modifying files"
    )
    parser.add_argument(
        "--no-backup", "-n",
        action="store_true", 
        help="Skip creating backup (not recommended)"
    )
    parser.add_argument(
        "--file", "-f",
        help="Translate a specific file instead of entire directory"
    )
    
    args = parser.parse_args()
    
    translator = LaTeXTitleTranslator()
    
    try:
        if args.file:
            # Single file mode
            if not os.path.isfile(args.file):
                print(f"Error: File does not exist: {args.file}")
                return 1
            
            print(f"Translating single file: {args.file}")
            if args.dry_run:
                print("DRY RUN MODE - No files will be modified")
            
            changed = translator.translate_file(args.file, args.dry_run)
            if changed:
                print(f"\n✓ File processed successfully")
            else:
                print(f"\n- No changes needed")
                
        else:
            # Directory mode
            print(f"Translating LaTeX files in: {args.directory}")
            if args.dry_run:
                print("DRY RUN MODE - No files will be modified")
            
            summary = translator.translate_directory(
                args.directory, 
                args.dry_run, 
                not args.no_backup
            )
            
            print(f"\n" + "="*50)
            print("TRANSLATION SUMMARY")
            print("="*50)
            print(f"Files processed: {summary['files_processed']}")
            print(f"Files changed: {summary['files_changed']}")
            print(f"Total translations: {summary['total_translations']}")
            
            if summary['files_changed'] > 0:
                if args.dry_run:
                    print("\nTo apply these changes, run the command without --dry-run")
                else:
                    print("\n✓ All translations completed successfully!")
            else:
                print("\n- No files needed translation")
        
        return 0
        
    except Exception as e:
        print(f"Error: {e}")
        return 1

if __name__ == "__main__":
    exit(main())